import platgo as pg


class GM(pg.Algorithm):
    """
    梯度法求解无约束优化问题： min(f(x))
    problem需可微，并且有相应计算梯度的函数
    """

    type: dict = {'single': True, 'multi': False, 'many': False, 'real': True, 'binary': False, 'permutation': False,
                  "large": True, 'expensive': False, 'constrained': False, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': True}

    def __init__(self, maxgen: int, problem: pg.Problem, beta: float = 0.5, sigma: float = 0.4) -> None:
        super().__init__(problem=problem, maxgen=maxgen)
        self.name = "GM"
        self.gk = None
        self.beta = beta
        self.step_len = pg.operators.ArmSearch(problem=problem, beta=beta, sigma=sigma)

    def go(self, N: int = 1, population: pg.Population = None) -> pg.Population:
        """
         if population is None, generate a new population with population size
        :param N: population size
        :param population: population to be optimized
        :return:
        """
        assert N or population, "N and population can't be both None"
        # 数学规划方法应该得到一个解而不是一个种群
        if population is None:
            pop = self.problem.init_pop(N=1)
        else:
            pop = population[0]
            self.problem.N = pop.decs.shape[0]
        self.problem.cal_obj(pop)

        while self.not_terminal(pop):
            self.gk = self.problem.g_fun(pop)  # 计算梯度
            dk = -self.gk  # 计算搜索方向
            mk = self.step_len(pop=pop, dk=dk, gk=self.gk)  # 用Armijo搜索求步长
            pop.decs = pop.decs + self.beta ** mk * dk
            self.problem.fix_decs(pop)
            self.problem.cal_obj(pop)
        return pop

